#!/usr/bin/env python
# -*- coding: utf-8 -*-

import numpy as np
from gnuradio import gr
from gnuradio import uhd
from gnuradio import blocks
from gnuradio import digital
from gnuradio import filter
import time
import os
import threading
import matplotlib.pyplot as plt
import glob
from scipy import linalg


class BidirectionalChannelEstimation(gr.top_block):
    def __init__(self):
        gr.top_block.__init__(self, "双向信道估计系统")

        # 硬件配置参数
        self.samp_rate = 3e5  # 优化后的采样率
        self.center_freq = 915e6  # 中心频率
        self.gain = 0  # 增益
        self.sps = 4  # 每符号采样数
        self.pilot_length = 127*2  # 导频序列长度
        self.num_iterations = 20  # 探测次数
        self.interval = 0.1  # 信道探测间隔时间
        self.current_role = "tx_rx"  # 当前角色
        self._running = False  # 运行状态标志
        self._tx_delay = 0.01  # 发射延迟
        self._buffer_size = 512
        # CSI数据存储
        self.csi_results = {
            'dir1': {'time': [], 'freq': []},  # 方向1的CSI数据
            'dir2': {'time': [], 'freq': []}  # 方向2的CSI数据
        }

        # 生成PN序列(伪随机噪声序列)
        self.glfsr = digital.glfsr(mask=0x83, seed=0xAC)
        self.pn_sequence = [self.glfsr.next_bit() for _ in range(self.pilot_length)]
        # 将PN序列转换为复数信号(BPSK调制)
        self.ref_signal = np.array([1.0 if bit else -1.0 for bit in self.pn_sequence],
                                   dtype=np.complex64)

        # 准备最小二乘估计矩阵
        # 使用更稳定的伪逆计算
        self.ref_matrix = linalg.circulant(self.ref_signal)
        self.pinv_ref_matrix = np.linalg.pinv(self.ref_matrix)

        # 初始化流图组件
        self.pilot_src = blocks.vector_source_b(self.pn_sequence * 20, True)  # 增加重复次数
        self.chunks_to_symbols = digital.chunks_to_symbols_bc([1.0 + 0j, -1.0 + 0j])  # 符号映射
        # 根升余弦滤波器
        self.rrc_filter = filter.interp_fir_filter_ccf(
            self.sps,
            filter.firdes.root_raised_cosine(1.0, self.sps, 1.0, 0.35, 41)
        )
        self.rx_sink = blocks.vector_sink_c()  # 接收数据存储

        # 初始化USRP设备(快速切换优化)
        self.init_usrp_fast_switch()

        # 创建结果保存目录
        os.makedirs("csi_results", exist_ok=True)

    def init_usrp_fast_switch(self):
        """初始化USRP设备(优化快速切换)"""
        # 设备1 (TX/RX)
        self.usrp1_tx = uhd.usrp_sink(
            device_addr="serial=31A92E2",
            stream_args=uhd.stream_args(
                cpu_format="fc32",
                args=f"send_frame_size={self._buffer_size}",  # 缓冲区大小
                channels=[0]
            ),
            antenna="TX/RX"
        )
        self.usrp1_rx = uhd.usrp_source(
            device_addr="serial=31A92E2",
            stream_args=uhd.stream_args(
                cpu_format="fc32",
                args=f"send_frame_size={self._buffer_size}",  # 缓冲区大小
                channels=[0]
            ),
            antenna="RX2"
        )

        # 设备2 (TX/RX)
        self.usrp2_tx = uhd.usrp_sink(
            device_addr="serial=31AC8A5",
            stream_args=uhd.stream_args(
                cpu_format="fc32",
                args=f"send_frame_size={self._buffer_size}",  # 缓冲区大小
                channels=[0]
            ),
            antenna="TX/RX"
        )
        self.usrp2_rx = uhd.usrp_source(
            device_addr="serial=31AC8A5",
            stream_args=uhd.stream_args(
                cpu_format="fc32",
                args=f"send_frame_size={self._buffer_size}",  # 缓冲区大小
                channels=[0]
            ),
            antenna="RX2"
        )

        # 配置通用参数(快速调谐优化)
        for usrp in [self.usrp1_tx, self.usrp1_rx, self.usrp2_tx, self.usrp2_rx]:
            usrp.set_samp_rate(self.samp_rate)
            usrp.set_center_freq(uhd.tune_request(self.center_freq))
            usrp.set_gain(self.gain)
            usrp.set_time_now(uhd.time_spec(0.0))
            usrp.set_command_time(uhd.time_spec(0.0))  # 启用定时命令

    def setup_flowgraph(self):
        """超快速流图配置(<2ms)"""
        self.lock()
        try:
            self.disconnect_all()

            if self.current_role == "tx_rx":
                # USRP1作为发射机，USRP2作为接收机
                self.connect((self.pilot_src, 0), (self.chunks_to_symbols, 0))
                self.connect((self.chunks_to_symbols, 0), (self.rrc_filter, 0))
                self.connect((self.rrc_filter, 0), (self.usrp1_tx, 0))
                self.connect((self.usrp2_rx, 0), (self.rx_sink, 0))
            else:
                # USRP2作为发射机，USRP1作为接收机
                self.connect((self.pilot_src, 0), (self.chunks_to_symbols, 0))
                self.connect((self.chunks_to_symbols, 0), (self.rrc_filter, 0))
                self.connect((self.rrc_filter, 0), (self.usrp2_tx, 0))
                self.connect((self.usrp1_rx, 0), (self.rx_sink, 0))

            self.rx_sink.reset()
        finally:
            self.unlock()

    def switch_roles(self):
        """立即切换收发角色（无延迟）"""
        if self.current_role == "tx_rx":
            # USRP1立即停止发射，USRP2立即停止接收
            self.usrp1_tx.set_stream_args(uhd.stream_args("fc32", args="flush=1"))
            self.usrp2_rx.set_stream_args(uhd.stream_args("fc32", args="flush=1"))
            self.current_role = "rx_tx"
        else:
            # USRP2立即停止发射，USRP1立即停止接收
            self.usrp2_tx.set_stream_args(uhd.stream_args("fc32", args="flush=1"))
            self.usrp1_rx.set_stream_args(uhd.stream_args("fc32", args="flush=1"))
            self.current_role = "tx_rx"

        # 立即重建流图连接
        self.setup_flowgraph()
        time.sleep(0.001)


    def estimate_channel(self, direction):
        # 设置精确的发射时间
        tx_time = uhd.time_spec(self._tx_delay)
        if direction == "dir1":
            self.usrp1_tx.set_start_time(tx_time)
        else:
            self.usrp2_tx.set_start_time(tx_time)

        # 开始采集
        self.rx_sink.reset()
        self.start()

        # 优化的采集窗口
        time.sleep(self._tx_delay + 0.05)

        self.stop()
        self.wait()



        # 处理数据
        rx_data = np.array(self.rx_sink.data())
        if len(rx_data) < self.pilot_length:
            print(f"警告: {direction}方向数据不足(收到{len(rx_data)}个样本)")
            return None, None

        valid_data = rx_data[:min(len(rx_data), self._buffer_size)]
        # 最小二乘估计
        y = valid_data[:self.pilot_length]
        h_est = np.dot(self.pinv_ref_matrix, y)
        csi = np.fft.fft(h_est, n=1024)

        # 存储结果
        self.csi_results[direction]['time'].append(h_est)
        self.csi_results[direction]['freq'].append(csi)

        return h_est, csi

    def save_results(self):
        """保存所有缓冲的结果"""
        timestamp = time.strftime("%Y%m%d_%H%M%S")
        for direction in ['dir1', 'dir2']:
            if len(self.csi_results[direction]['time']) > 0:
                filename = f"csi_results/{direction}_csi_{timestamp}.npz"
                np.savez(
                    filename,
                    time_response=np.array(self.csi_results[direction]['time']),
                    freq_response=np.array(self.csi_results[direction]['freq'])
                )
                print(f"已保存{len(self.csi_results[direction]['time'])}个{direction}结果")



    def run(self):
        """主运行循环"""
        self._running = True
        try:
            # 初始连接配置
            self.setup_flowgraph()

            for i in range(self.num_iterations):
                if not self._running:
                    break

                print(f"\n=== 第 {i + 1}/{self.num_iterations} 次迭代 ===")

                # 方向1测量
                print("[USRP1→USRP2] 开始测量...")
                self.estimate_channel("dir1")

                # 切换方向
                self.switch_roles()

                # 方向2测量
                print("[USRP2→USRP1] 开始测量...")
                self.estimate_channel("dir2")

                # 恢复初始状态
                self.switch_roles()

                if i < self.num_iterations - 1:
                    time.sleep(self.interval)

            # 保存所有结果
            self.save_results()



        except KeyboardInterrupt:
            print("用户中断-保存已收集数据...")
            self.save_results()
        except Exception as e:
            print(f"运行时错误: {str(e)} - 保存已收集数据...")
            self.save_results()
        finally:
            self._running = False
            print("测量完成")


def find_latest_csi_files():
    """查找并返回最新的匹配CSI文件对路径

    返回:
        tuple: 返回最新的(dir1文件路径, dir2文件路径)，如果没有匹配对则返回(None, None)
    """
    try:
        # 获取所有CSI文件并按修改时间排序(最新的在前面)
        dir1_files = sorted(glob.glob(os.path.join("csi_results", "dir1_csi_*.npz")),
                            key=os.path.getmtime, reverse=True)
        dir2_files = sorted(glob.glob(os.path.join("csi_results", "dir2_csi_*.npz")),
                            key=os.path.getmtime, reverse=True)

        # 如果没有文件，直接返回None
        if not dir1_files or not dir2_files:
            print("未找到足够的CSI文件")
            return None, None

        # 创建一个字典来存储时间戳和文件路径的映射
        timestamp_map = {}

        # 处理dir1文件
        for file_path in dir1_files:
            try:
                # 从文件名提取时间戳部分，例如：dir1_csi_20230815_143022.npz
                filename = os.path.basename(file_path)
                timestamp = '_'.join(filename.split('_')[2:]).replace('.npz', '')
                timestamp_map[timestamp] = {'dir1': file_path}
            except Exception as e:
                print(f"处理文件{file_path}时出错: {str(e)}")
                continue

        # 处理dir2文件并尝试匹配
        for file_path in dir2_files:
            try:
                filename = os.path.basename(file_path)
                timestamp = '_'.join(filename.split('_')[2:]).replace('.npz', '')

                if timestamp in timestamp_map:
                    timestamp_map[timestamp]['dir2'] = file_path
                    # 找到第一个匹配的就返回(因为文件已经按时间排序)
                    return (timestamp_map[timestamp]['dir1'],
                            timestamp_map[timestamp]['dir2'])
            except Exception as e:
                print(f"处理文件{file_path}时出错: {str(e)}")
                continue

        print("未找到时间戳匹配的CSI文件对")
        return None, None

    except Exception as e:
        print(f"查找CSI文件时发生错误: {str(e)}")
        return None, None


def plot_csi_comparison( dir1_file, dir2_file):
    """绘制两个方向的CSI幅度对比图

    参数:
        dir1_file (str): 方向1的CSI文件路径
        dir2_file (str): 方向2的CSI文件路径
    """
    samp_rate = 1e6

    if dir1_file is None or dir2_file is None:
        print("错误: 必须提供两个有效的文件路径")
        return


    # 加载CSI数据
    dir1_data = np.load(dir1_file)
    dir2_data = np.load(dir2_file)

    # 获取时域信道
    # dir1_time = dir1_data['freq_response']
    # dir2_time = dir2_data['freq_response']
    dir1_time = dir1_data['time_response']
    dir2_time = dir2_data['time_response']

    # 计算幅度(dB)
    dir1_mag = 20 * np.log10(1e-16+np.abs(dir1_time.reshape((-1,))))  # 取第一次测量的数据
    dir2_mag = 20 * np.log10(1e-16+np.abs(dir2_time.reshape((-1,))))

    # 创建频率轴
    # freq_axis = np.linspace(-samp_rate / 2, samp_rate / 2, len(dir1_mag))

    # 绘制图形
    plt.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体
    plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题
    plt.figure(figsize=(12, 8))

    # 绘制两个方向的CSI幅度对比
    plt.plot( dir1_mag, label='USRP1→USRP2')
    plt.plot( dir2_mag, label='USRP2→USRP1', alpha=0.7)
    plt.title('CSI幅度对比')
    plt.xlabel('时间步')
    plt.ylabel('幅度 (dB)')
    plt.grid(True)
    plt.legend()

    # 计算并显示相似度
    correlation = np.corrcoef(dir1_mag, dir2_mag)[0, 1]
    plt.text(0.05, 0.95, f'相似度(相关系数): {correlation:.4f}',
             transform=plt.gca().transAxes, bbox=dict(facecolor='white', alpha=0.8))

    # 保存图形
    timestamp = time.strftime("%Y%m%d_%H%M%S")
    plot_filename = f"csi_results/csi_comparison_{timestamp}.png"
    plt.savefig(plot_filename)
    plt.show()
    # plt.close()

    print(f"CSI对比图已保存为: {plot_filename}")

    # 打印相似度结论
    if correlation > 0.8:
        print("结论: 两个方向的CSI幅度高度相似")
    elif correlation > 0.5:
        print("结论: 两个方向的CSI幅度中等相似")
    else:
        print("结论: 两个方向的CSI幅度差异较大")



def main():
    print("=== 快速切换双向信道估计系统 ===")
    print("初始化硬件...")

    try:
        estimator = BidirectionalChannelEstimation()
        print("系统准备就绪，开始测量...")
        estimator.run()
    except Exception as e:
        print(f"系统初始化失败: {str(e)}")

if __name__ == "__main__":
    main()

    # 获取最新的文件对
    dir1_path, dir2_path = find_latest_csi_files()
    # 绘制对比图
    if dir1_path and dir2_path:
        plot_csi_comparison(dir1_path, dir2_path)
    else:
        print("无法找到有效的CSI文件对")
